const Mesh = @import("./mesh/mesh.zig").Mesh;
const std = @import("std");
const zlbmath = @import("zlbmath");
const lbmath = zlbmath.lbmath;
pub fn KD_Node(comptime V: type, comptime T: type) type {
    return struct {
        const Self = @This();
        left: ?*Self = null,
        right: ?*Self = null,
        parent: ?*Self = null,
        dim: u32 = 0,
        value: V,
        prop: T,
        pub fn init(self: *Self) void {
            self.left = null;
            self.right = null;
            self.parent = null;
            self.dim = 0;
        }
        inline fn kd_brother(self: *const Self) ?*Self {
            if (self.parent == null) return null;
            if (self.parent.?.left == self) return self.parent.?.right;
            return self.parent.?.left;
        }
        inline fn is_intersecting_with_sphere(self: *const Self, r: f64, k: V) bool {
            const dist = @abs(k.point[self.dim] - self.value.point[self.dim]);
            return dist <= r;
        }
        inline fn swap(values: []V, left: u64, right: u64) void {
            const temp = values[left];
            values[left] = values[right];
            values[right] = temp;
        }
        fn sort_vertices(values: []V, dim: usize) void {
            if (values.len <= 1) {
                return;
            }
            // var d=dim%values[0].point_size;
            var m = 0;
            // for (int i=left+1;i<=right;i++)
            for (1..values.len) |i| {
                if (values[i].point[dim] < values[0].point[dim]) {
                    m += 1;
                    swap(values, m, i);
                }
            }
            swap(values, 0, m);
            sort_vertices(values[0..m], dim);
            sort_vertices(values[(m + 1)..], dim);
        }
        // --- 新增：分区辅助函数 (Lomuto Partition Scheme) ---
        // 作用：选取最后一个元素为轴点，重排数组，使得所有小于轴点的在左，大于等于的在右
        // 返回轴点最终所在的位置索引
        fn partition(values: []V, dim: usize) usize {
            if (values.len <= 1) return 0;
            var m = 0;
            // for (int i=left+1;i<=right;i++)
            for (1..values.len) |i| {
                if (values[i].point[dim] < values[0].point[dim]) {
                    m += 1;
                    swap(values, m, i);
                }
            }
            swap(values, 0, m);
            return m;
        }

        // --- 新增：Quickselect 算法，用于高效查找中位数 ---
        // 作用：在 O(n) 平均时间内找到数组的中位数，并将其移动到正确位置。
        // 返回中位数的索引。执行后，数组被分区，但不是完全排序。
        fn findMedianAndPartition(values: []V, dim: usize) usize {
            var low: usize = 0;
            var high: usize = values.len - 1;
            const median_idx: usize = values.len / 2;

            while (low < high) {
                // 对 low..high+1 的子切片进行分区
                // 注意：partition函数作用于整个子切片，所以需要调整偏移
                const pivot_idx = low + partition(values[low .. high + 1], dim);

                if (pivot_idx == median_idx) {
                    return median_idx; // 找到了中位数
                } else if (pivot_idx < median_idx) {
                    low = pivot_idx + 1; // 中位数在右侧
                } else {
                    high = pivot_idx - 1; // 中位数在左侧
                }
            }
            return median_idx;
        }
        // 移植的create_kd_node函数
        pub fn create_kd_node(allocator: std.mem.Allocator, values: []V, dim: usize) !?*Self {
            if (values.len == 0) {
                return null;
            }

            // 获取点的维度大小
            const point_size = values[0].point.len;

            // 对当前维度进行排序
            // sort_vertices(values, 0, values.len - 1, dim % point_size);
            const d = dim % point_size;
            const median_idx = findMedianAndPartition(values, d);

            // 步骤2: 创建新节点
            const node = try allocator.create(Self);
            node.* = .{ // 使用 struct literal 初始化，更简洁
                .left = null,
                .right = null,
                .parent = null,
                .dim = d,
                .value = values[median_idx], // 使用真正的中位数作为节点值
                .prop = .{}, // 假设 T 有默认值
            };
            // 如果只有一个元素，直接返回
            if (values.len == 1) {
                node.value = values[0];
                return node;
            }

            // 递归构建左右子树
            node.left = try create_kd_node(allocator, values[0..median_idx], dim + 1);
            node.right = try create_kd_node(allocator, values[(median_idx + 1)..], dim + 1);

            // 设置父节点
            if (node.left) |left| {
                left.parent = node;
            }
            if (node.right) |right| {
                right.parent = node;
            }

            return node;
        }
        pub fn find_neareast_sphere(self: *const Self, v: V, r: f64, list: *std.ArrayList(V), allocator: std.mem.Allocator) void {
            const dis = v.point[self.dim] - self.value.point[self.dim];
            if (dis > r) {
                if (self.right) |ri| {
                    ri.find_neareast_sphere(v, r, list, allocator);
                }
            } else if (dis < -r) {
                if (self.left) |le| {
                    le.find_neareast_sphere(v, r, list, allocator);
                }
            } else {
                var dis1 = (v.point[0] - self.value.point[0]) * (v.point[0] - self.value.point[0]);
                for (v.point[1..], self.value.point[1..]) |p1, p2| {
                    dis1 += (p1 - p2) * (p1 - p2);
                }
                // if (lbmath.distance_of_two_points(v.point, self.value.point) <= r)
                if (dis1 <= r * r) {
                    list.append(allocator, self.value) catch unreachable;
                }
                if (self.right) |ri| {
                    ri.find_neareast_sphere(v, r, list, allocator);
                }
                if (self.left) |le| {
                    le.find_neareast_sphere(v, r, list, allocator);
                }
            }
        }
        fn dis_sq_two_distance(a: anytype, b: []const (@TypeOf(a[0]))) @TypeOf(a[0]) {
            var re: @TypeOf(a[0]) = 0;
            for (a, b) |va, vb| {
                re += (va - vb) * (va - vb);
            }
            return re;
        }
        // KNN 需要一个上下文来比较堆中元素的距离
        fn greaterThan(context: V, a: V, b: V) std.math.Order {
            var dis1 = (a.point[0] - context.point[0]) * (a.point[0] - context.point[0]);
            var dis2 = (b.point[0] - context.point[0]) * (b.point[0] - context.point[0]);
            for (1..context.point.len) |i| {
                dis1 += (a.point[i] - context.point[i]) * (a.point[i] - context.point[i]);
                dis2 += (b.point[i] - context.point[i]) * (b.point[i] - context.point[i]);
            }
            return if (dis1 < dis2) .gt else .lt;
        }
        // 按照优先顺序递归计算邻近点
        // 后一个根据情况剪枝
        fn priority_tool_knn(
            self: *const Self,
            target: V,
            k: usize,
            queue: *std.PriorityQueue(V, V, greaterThan),
            dis: anytype,
        ) void {
            if (dis < 0) {
                if (self.left) |le| le.knn_search_recursive(target, k, queue);

                if (dis * dis < dis_sq_two_distance(queue.items[0].point, target.point)) {
                    if (self.right) |ri| ri.knn_search_recursive(target, k, queue);
                }
            } else {
                if (self.right) |ri| ri.knn_search_recursive(target, k, queue);
                if (dis * dis < dis_sq_two_distance(queue.items[0].point, target.point)) {
                    if (self.left) |le| le.knn_search_recursive(target, k, queue);
                }
            }
        }
        // 递归辅助函数
        fn knn_search_recursive(
            self: *const Self,
            target: V,
            k: usize,
            queue: *std.PriorityQueue(V, V, greaterThan),
        ) void {
            const dis = target.point[self.dim] - self.value.point[self.dim];

            if (queue.items.len < k) {
                queue.add(self.value);
                if (dis < 0) {
                    if (self.left) |le| le.knn_search_recursive(target, k, queue);
                    if (queue.items.len < k or dis * dis < dis_sq_two_distance(queue.items[0].point, target.point)) {
                        if (self.right) |ri| ri.knn_search_recursive(target, k, queue);
                    }
                } else {
                    if (self.right) |ri| ri.knn_search_recursive(target, k, queue);
                    if (queue.items.len < k or dis * dis < dis_sq_two_distance(queue.items[0].point, target.point)) {
                        if (self.left) |le| le.knn_search_recursive(target, k, queue);
                    }
                }
            } else {
                const dis1 = dis_sq_two_distance(queue.items[0].point, target.point);
                const dis2 = dis_sq_two_distance(self.value.point, target.point);
                if (dis2 < dis1) {
                    queue.remove();
                    queue.add(self.value);
                    self.priority_tool_knn(target, k, queue, dis);
                } else {
                    if (dis * dis > dis1) {
                        if (dis > 0) {
                            // dis*dis>dis1 且dis>0说明左侧节点距离目标点大于dis1,故而剪枝
                            // 下面同理
                            if (self.right) |ri| ri.knn_search_recursive(target, k, queue);
                        } else {
                            if (self.left) |le| le.knn_search_recursive(target, k, queue);
                        }
                    } else {
                        self.priority_tool_knn(target, k, queue, dis);
                    }
                }
            }
        }

        // 公开的调用接口
        pub fn find_nearest_k_vertices(
            self: *const Self,
            target: V,
            k: usize,
            allocator: std.mem.Allocator,
        ) ?[]V {
            if (k == 0) return null;

            var queue = std.PriorityQueue(V, V, greaterThan).init(allocator, target);

            // 确保退出时释放内存
            defer queue.deinit();

            // 3. 执行递归搜索
            self.knn_search_recursive(target, k, &queue);
            if (queue.len <= 0) return null;
            const re = allocator.alloc(V, queue.items.len) catch return null;
            for (0..re.len) |i| {
                re[i] = queue.remove();
            }
            // @memmove(re, queue.items);
            return re;
        }
        pub fn find_nearest_vertex(
            self: *const Self,
            target: V,
        ) ?V {
            var item: [1]V = undefined;
            var queue = std.PriorityQueue(V, V, greaterThan){
                .items = item[0..0],
                .cap = 1,
                .allocator = std.heap.page_allocator,
                .context = target,
            };

            self.knn_search_recursive(target, 1, &queue);
            return if (queue.items.len > 0) queue.items[0] else null;
        }
    };
}
// test "pri_deque"{
//     std.PriorityQueue(u32, void, lessThan);
// }
// fn insertion_sort(values: []V, dim: u32) void {
//     // ... (这个函数保持不变，是正确的)
//     var i: usize = 1;
//     while (i < values.len) : (i += 1) {
//         const current = values[i];
//         var j = i;
//         while (j > 0 and values[j - 1].point[dim] > current.point[dim]) {
//             values[j] = values[j - 1];
//             j -= 1;
//         }
//         values[j] = current;
//     }
// }
// // --- 最终优化版：使用3路快排，处理重复元素 ---
// pub fn robust_sort_vertices(values: []V, dim: u32) void {
//     // 优化1: 小数组切换为插入排序
//     const INSERTION_SORT_CUTOFF = 16;
//     if (values.len < INSERTION_SORT_CUTOFF) {
//         insertion_sort(values, dim);
//         return;
//     }

//     // 优化2: 使用三数取中法获得一个好的轴点 *值*
//     // (这里我们只取值，分区时再移动元素)
//     const mid_idx = values.len / 2;
//     const high_idx = values.len - 1;
//     var pivot_value: @TypeOf(values[0].point[0]);
//     // 简单的比较逻辑来找出三个候选的中位数
//     const a = values[0].point[dim];
//     const b = values[mid_idx].point[dim];
//     const c = values[high_idx].point[dim];
//     if ((a > b) != (a > c)) { // a is the median
//         pivot_value = a;
//     } else if ((b > a) != (b > c)) { // b is the median
//         pivot_value = b;
//     } else { // c is the median
//         pivot_value = c;
//     }

//     // 优化3: 3路分区 (荷兰国旗问题算法)
//     // |  < pivot  |  == pivot  |  未处理  |  > pivot  |
//     // |-----------|------------|----------|-----------|
//     // 0          lt           i          gt        len
//     var lt: usize = 0; // "小于"区域的右边界 (不含)
//     var i: usize = 0; // 当前处理元素的指针
//     var gt: usize = values.len; // "大于"区域的左边界 (不含)

//     while (i < gt) {
//         const current_val = values[i].point[dim];
//         if (current_val < pivot_value) {
//             swap(values, lt, i);
//             lt += 1;
//             i += 1;
//         } else if (current_val > pivot_value) {
//             gt -= 1;
//             swap(values, i, gt);
//         } else { // current_val == pivot_value
//             i += 1;
//         }
//     }

//     // 递归地对左右两个子数组进行排序
//     // 注意：中间等于 pivot_value 的部分 [lt..gt] 无需再排序
//     robust_sort_vertices(values[0..lt], dim);
//     robust_sort_vertices(values[gt..], dim);
// }
// --- K 最近邻搜索 (KNN) ---
